import torch
from configparser import ConfigParser
from transformers import BertTokenizer, BertModel
from sklearn.metrics.pairwise import cosine_similarity


# NOTE:
# 可用, 基于bert-base-chinese
# 加载时间笔记本十几秒吧


class TextScoreBert:
    _instance = None
    _PATH_MODEL = None
    _tokenizer = None
    _model = None

    def __init__(self, path_config):
        if self._instance is not None:
            return
        # 创建ConfigParser对象
        config = ConfigParser()
        # 读取配置文件
        config.read(path_config, encoding='utf-8')
        # 获取default部分的api_key
        self._PATH_MODEL = config.get('default', 'bert_base_chinese')
        print("bert_base_chinese from:", self._PATH_MODEL)
        if self._tokenizer == None or self._model == None:
            # 指定已存在的翻译模型和tokenizer
            self._tokenizer = BertTokenizer.from_pretrained(self._PATH_MODEL)
            self._model = BertModel.from_pretrained(self._PATH_MODEL)

    def __new__(cls, path_config):
        # print("__new__")
        if cls._instance is None:
            instance = super().__new__(cls)
            instance.__init__(path_config)  # 如果需要初始化操作
            cls._instance = instance
        return cls._instance

    def get_embedding(self, text):
        # 编码输入文本
        inputs = self._tokenizer(text, return_tensors='pt')
        # 获取文本嵌入
        outputs = self._model(**inputs)
        return outputs.last_hidden_state.mean(dim=1).detach().numpy()

    def similar_text(self, text1, text2):
        # 获取文本嵌入
        emb1 = self.get_embedding(text1)
        emb2 = self.get_embedding(text2)
        # 计算余弦相似度
        similarity = cosine_similarity(emb1, emb2)[0][0]
        return similarity


def test():
    # 使用示例
    llm = TextScoreBert("../files/config.inf")
    text1 = "这是第一个文本。"
    text2 = "这是第二个文本，它有一些不同。"
    similar_val = llm.similar_text(text1, text2)
    print(f"文本相似度: {similar_val}")
    print("----------------------------------------------------------------------------------------------")
    text1 = f"这次通过VR虚拟实验室进行的一系列关于反相器的实验，体验了从原理设计，版图设计，芯片制作，最后到芯片封装测试一系列完整的流程，了解到了集成芯片制作的总体思路和具体技术，收获颇深。"
    text2 = f"完成反相器制造，学习集成电路设计制造、测试和流程。"
    similar_val = llm.similar_text(text1, text2)
    print(f"文本相似度: {similar_val}")




if __name__=="__main__":
    test()